package com.googlecode.quillen.repository.aws;

import static com.googlecode.quillen.util.Utils.logDebug;
import com.googlecode.quillen.util.ResultConsumer;
import com.googlecode.quillen.util.Pair;
import com.googlecode.quillen.domain.Attributes;
import com.googlecode.quillen.domain.AttributeStorageException;
import com.googlecode.quillen.repository.AttributeStorage;
import com.googlecode.quillen.repository.AttributeStorageQuery;
import com.xerox.amazonws.sdb.*;
import com.xerox.amazonws.common.AWSError;

import java.util.*;
import java.util.regex.Pattern;
import java.util.regex.Matcher;

import org.apache.commons.lang.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * Created by IntelliJ IDEA.
 * User: greg
 * Date: Dec 8, 2008
 * Time: 10:26:29 PM
 */
public class SDBAttributeStorageImpl implements AttributeStorage
{
    private static final Log LOG = LogFactory.getLog(SDBAttributeStorageImpl.class);

    private static class QueryPrinter
    {
        private final String database;

        private QueryPrinter(String database)
        {
            this.database = database;
        }

        public String print(AttributeStorageQuery.Condition condition)
        {
            if (condition instanceof AttributeStorageQuery.ComplexCondition)
            {
                AttributeStorageQuery.ComplexCondition c = (AttributeStorageQuery.ComplexCondition) condition;
                return "((" + print(c.left) + ") " + c.conjunction + " (" + print(c.right) + "))";
            }
            else if (condition instanceof AttributeStorageQuery.LeafCondition)
            {
                AttributeStorageQuery.LeafCondition c = (AttributeStorageQuery.LeafCondition) condition;
                return escapeNameForQuery(c.attribute) + " " + c.operator + " " + escapeValueForQuery(c.value);
            }
            else if (condition instanceof AttributeStorageQuery.NullCondition)
            {
                return escapeNameForQuery(((AttributeStorageQuery.NullCondition) condition).attribute) + " is null";
            }
            else if (condition instanceof AttributeStorageQuery.NotNullCondition)
            {
                return escapeNameForQuery(((AttributeStorageQuery.NotNullCondition) condition).attribute) + " is not null";
            }
            else
            {
                throw new IllegalStateException("unknown condition type: " + condition.getClass());
            }
        }

        public String print(AttributeStorageQuery query)
        {
            if (query != null)
            {
                String result = "select " + (query.getProjection().isEmpty() ? "*" : StringUtils.join(query.getProjection(), ", ")) + " " +
                                "from " + database;

                if (query.getCondition() != null)
                {
                    result += " where " + print(query.getCondition());
                }

                // TODO: order by

                return result;
            }
            else
            {
                return null;
            }
        }
    }

    private final SimpleDB simpleDb;

    public SDBAttributeStorageImpl(SimpleDB simpleDb)
    {
        this.simpleDb = simpleDb;
        simpleDb.setMaxRetries(10);
    }

    public void createDatabase(String database) throws AttributeStorageException
    {
        try
        {
            simpleDb.createDomain(database);
        }
        catch (SDBException e)
        {
            throw wrapException(e);
        }
    }

    public void deleteDatabase(String database) throws AttributeStorageException
    {
        try
        {
            simpleDb.deleteDomain(database);
        }
        catch (SDBException e)
        {
            throw wrapException(e);
        }
    }

    public Attributes get(String database, String key, String... attributeNames) throws AttributeStorageException
    {
        List<String> attributesToGet = null;

        if (attributeNames != null && attributeNames.length > 0)
        {
            attributesToGet = Arrays.asList(attributeNames);
        }

        try
        {
            Attributes attrs = toAttributes(simpleDb.getDomain(database).getItem(key).getAttributes(attributesToGet));

            return attrs.isEmpty() ? null : attrs;
        }
        catch (SDBException e)
        {
            throw wrapException(e);
        }
    }

    public void put(final String database, final String key, final Attributes attributes) throws AttributeStorageException
    {
        try
        {
            Item item = simpleDb.getDomain(database).getItem(key);
            List<ItemAttribute> itemAttributes = new ArrayList<ItemAttribute>();
            Attributes attrsCopy = attributes.copy();

            for (String k : attrsCopy.getKeys())
            {
                for (String v : attrsCopy.get(k))
                {
                    itemAttributes.add(new ItemAttribute(k, v, attributes.isReplace(k)));

                    if (itemAttributes.size() == 255)
                    {
                        putAndRemoveSuccessful(item, itemAttributes, attributes);
                        itemAttributes.clear();
                    }
                }
            }

            if (!itemAttributes.isEmpty())
            {
                putAndRemoveSuccessful(item, itemAttributes, attributes);
            }
        }
        catch (SDBException e)
        {
            throw wrapException(e);
        }
    }

    public void batchPut(final String database, final Map<String, Attributes> attributes) throws AttributeStorageException
    {
        try
        {
            Domain domain = simpleDb.getDomain(database);
            Map<String, List<ItemAttribute>> batch = new HashMap<String, List<ItemAttribute>>();

            for (Map.Entry<String, Attributes> entry : new HashMap<String, Attributes>(attributes).entrySet())
            {
                final String key = entry.getKey();
                final Attributes attrs = entry.getValue();

                // Consider higher priority and replaceable attributes first so they have a higher likelihood of going into the batch
                //
                List<String> sortedKeys = new ArrayList<String>(attrs.getKeys());
                Collections.sort(sortedKeys, new Comparator<String>()
                {
                    public int compare(String left, String right)
                    {
                        int result = ((Integer) attrs.getPriority(right)).compareTo(attrs.getPriority(left));
                        if (result != 0)
                        {
                            return result;
                        }

                        result = ((Boolean) attrs.isReplace(right)).compareTo(attrs.isReplace(left));
                        if (result != 0)
                        {
                            return result;
                        }

                        return left.compareTo(right);
                    }
                });

                // Build a list of up to 256 ItemAttributes for this item key
                //
                List<ItemAttribute> itemAttributes = new ArrayList<ItemAttribute>();
                fillItemAttributes:
                for (String k : sortedKeys)
                {
                    for (String v : attrs.get(k))
                    {
                        if (itemAttributes.size() < 256)
                        {
                            itemAttributes.add(new ItemAttribute(k, v, attrs.isReplace(k)));
                        }
                        else
                        {
                            break fillItemAttributes;
                        }
                    }
                }
                batch.put(key, itemAttributes);

                // If we have a batch of 25, go ahead and do a BatchPut
                //
                if (batch.size() == 25)
                {
                    batchPutAndRemoveSuccessful(domain, batch, attributes);
                    batch.clear();
                }
            }

            batchPutAndRemoveSuccessful(domain, batch, attributes);
        }
        catch (SDBException e)
        {
            throw wrapException(e);
        }
    }

    private static Pattern NUMBER_ATTRS_EXCEEDED_PATTERN = Pattern.compile(".*Too many attributes in Items? (.+)\\.");

    // TODO: need to worry about 1MB request size limit?
    private void batchPutAndRemoveSuccessful(Domain domain, Map<String, List<ItemAttribute>> batch, Map<String, Attributes> attributes) throws SDBException
    {
        boolean retry;
        do
        {
            retry = false;

            for (Map.Entry<String, List<ItemAttribute>> entry : new HashMap<String, List<ItemAttribute>>(batch).entrySet())
            {
                if (entry.getValue().isEmpty())
                {
                    batch.remove(entry.getKey());
                }
            }

            if (!batch.isEmpty())
            {
                try
                {
                    domain.batchPutAttributes(batch);
                }
                catch (SDBException e)
                {
                    if (e.getErrors() != null)
                    {
                        for (AWSError error : e.getErrors())
                        {
                            if ("NumberItemAttributesExceeded".equals(error.getCode()))
                            {
                                // We have to parse the error mesage to know which items were at fault
                                //
                                Matcher matcher = NUMBER_ATTRS_EXCEEDED_PATTERN.matcher(error.getMessage());
                                if (matcher.matches())
                                {
                                    for (String item : StringUtils.splitByWholeSeparator(matcher.group(1), ", "))
                                    {
                                        logDebug(LOG, "too many attributes in item %s.  removing from batch and retrying.", item);
                                        batch.remove(item);
                                    }
                                }

                                retry = true;
                                break;
                            }
                        }
                    }

                    if (!retry)
                    {
                        throw e;
                    }
                }
            }
        }
        while (retry);

        if (attributes != null)
        {
            for (Map.Entry<String, List<ItemAttribute>> entry : batch.entrySet())
            {
                String key = entry.getKey();
                Attributes attrs = attributes.get(key);

                for (ItemAttribute itemAttr : entry.getValue())
                {
                    attrs.remove(itemAttr.getName(), itemAttr.getValue());
                }

                if (attrs.isEmpty())
                {
                    attributes.remove(key);
                }
            }
        }
    }

    private void putAndRemoveSuccessful(Item item, List<ItemAttribute> itemAttributes, Attributes attributes) throws SDBException
    {
        try
        {
            item.putAttributes(itemAttributes);
        }
        catch (SDBException e)
        {
            // TODO: errors are on inner exception?
//            if (e.getErrors() != null)
//            {
//                for (AWSError error : e.getErrors())
//                {
//                    if ("NumberItemAttributesExceeded".equals(error.getCode()))
//                    {
//                        return;
//                    }
//                }
//            }

            if (StringUtils.contains(e.getMessage(), "Too many attributes in Item"))
            {
                return;
            }

            throw e;
        }

        for (ItemAttribute itemAttr : itemAttributes)
        {
            attributes.remove(itemAttr.getName(), itemAttr.getValue());
        }
    }

    public void delete(String database, String key) throws AttributeStorageException
    {
        try
        {
            simpleDb.getDomain(database).deleteItem(key);
        }
        catch (SDBException e)
        {
            throw wrapException(e);
        }
    }

    public void delete(String database, String key, Attributes attributes) throws AttributeStorageException
    {
        try
        {
            Item item = simpleDb.getDomain(database).getItem(key);
            List<ItemAttribute> itemAttrs = new ArrayList<ItemAttribute>();

            for (String k : attributes.getKeys())
            {
                for (String v : attributes.get(k))
                {
                    itemAttrs.add(new ItemAttribute(k, v, false));

                    if (itemAttrs.size() == 255)
                    {
                        item.deleteAttributes(itemAttrs);
                        itemAttrs.clear();
                    }
                }
            }

            if (!itemAttrs.isEmpty())
            {
                item.deleteAttributes(itemAttrs);
            }
        }
        catch (SDBException e)
        {
            throw wrapException(e);
        }
    }

    public void list(String database, ResultConsumer<Pair<String, Attributes>> consumer) throws AttributeStorageException
    {
        find(database, null, consumer);
    }

    public void find(String database, AttributeStorageQuery query, ResultConsumer<Pair<String, Attributes>> consumer) throws AttributeStorageException
    {
        String queryString = new QueryPrinter(database).print(query);

        try
        {
            QueryWithAttributesResult result;
            String nextToken = null;
            do
            {
                if (StringUtils.isNotEmpty(queryString))
                {
                    result = simpleDb.getDomain(database).selectItems(queryString, nextToken);
                }
                else
                {
                    result = simpleDb.getDomain(database).listItemsWithAttributes(null, null, nextToken, 0);
                }

                for (Map.Entry<String, List<ItemAttribute>> entry : result.getItems().entrySet())
                {
                    if (consumer != null)
                    {
                        consumer.newResult(Pair.get(entry.getKey(), toAttributes(entry.getValue())));
                    }
                }
            }
            while ((nextToken = result.getNextToken()) != null);
        }
        catch (Exception e)
        {
            throw wrapException(e);
        }
    }

    private Attributes toAttributes(List<ItemAttribute> itemAttributes)
    {
        Attributes result = new Attributes();

        for (ItemAttribute attr : itemAttributes)
        {
            result.add(attr.getName(), attr.getValue());
        }

        return result;
    }

    @SuppressWarnings("ThrowableInstanceNeverThrown")
    private AttributeStorageException wrapException(Exception e)
    {
        if (e instanceof AttributeStorageException)
        {
            return (AttributeStorageException) e;
        }
        else
        {
            return new AttributeStorageException(e);
        }
    }

    private static String escapeValueForQuery(String str)
    {
        return "'" + StringUtils.replaceEach(str, new String[] {"'", "\""}, new String[] {"''", "\"\""}) + "'";
    }

    private static String escapeNameForQuery(String str)
    {
        if (!"itemName()".equals(str))
        {
            return "`" + StringUtils.replace(str, "`", "``") + "`";
        }
        else
        {
            return str;
        }
    }
}
